# Edge Impulse - OpenMV Object Detection Example

import sensor, image, time, os, tf, math, uos, gc
import lcd

sensor.reset()                         # Reset and initialize the sensor.
sensor.set_pixformat(sensor.RGB565)    # Set pixel format to RGB565 (or GRAYSCALE)
sensor.set_framesize(sensor.QQVGA2)      # Set frame size to QQVGA2 (128x160)
#sensor.set_windowing((240, 240))       # Set 240x240 window.
sensor.skip_frames(time=2000)          # Let the camera adjust.
lcd.init() # Initialize the lcd screen.

net = None              #要导入的模型文件，先设置为空
labels = None           #要导入的标签文件，先设置为空
min_confidence = 0.5    #最小置信度 0.5
Mode_Flag = 1
#模式1：开机识别数字
#模式2：巡线
#模式3：再次识别数字，发送左转or右转指令


try:
    # load the model, alloc the model file on the heap if we have at least 64K free after loading
    net = tf.load("trained.tflite", load_to_fb=uos.stat('trained.tflite')[6] > (gc.mem_free() - (64*1024)))
except Exception as e:      #加载异常的报错Failed to load
    raise Exception('Failed to load "trained.tflite", did you copy the .tflite and labels.txt file onto the mass-storage device? (' + str(e) + ')')

try:
    # 对列表labels进行赋值，列表内容是labels.txt内每一行的内容，以换行符分割元素
    labels = [line.rstrip('\n') for line in open("labels.txt")]
except Exception as e:
    raise Exception('Failed to load "labels.txt", did you copy the .tflite and labels.txt file onto the mass-storage device? (' + str(e) + ')')

colors = [ # Add more colors if you are detecting more than 7 types of classes at once.
    (  0,   0,   0),    #白色 background
    (255,   0,   0),    #红色 1
    (  0, 255,   0),    #绿色 2
    (  0,   0, 255),    #蓝色 3
    (255, 255,   0),    #黄色 4
    (255,   0, 255),    #品红 5
    (  0, 255, 255),    #青色 6
    (255, 255, 255),
    (230,  25, 180),    #多写两行，解决报错 ErrorIndex
]

clock = time.clock()
while(True):
    clock.tick()

    img = sensor.snapshot()
    img.draw_line((64, 0, 64, 160), color = 255)        #划屏幕中线
    lcd.display(img) # Take a picture and display the image.
    # detect() returns all objects found in the image (splitted out per class already)
    # we skip class index 0, as that is the background, and then draw circles of the center
    # of our objects

    # 1.串口接收数据，处理 Mode_Flag 标志位


    if (Mode_Flag == 1):    #如果是模式1，开机识别数字


    if (Mode_Flag == 3):    #如果是模式1

        for i, detection_list in enumerate(net.detect(img, thresholds=[(math.ceil(min_confidence * 255), 255)])):
            if (i == 0): continue # background class
            if (len(detection_list) == 0): continue # no detections for this class?

            print("********** %s **********" % labels[i])
            for d in detection_list:
                [x, y, w, h] = d.rect()
                center_x = math.floor(x + (w / 2))
                center_y = math.floor(y + (h / 2))
                print('x %d\ty %d' % (center_x, center_y))

            if (center_x <= 64):    # 中心坐标x <= 64,识别的图像在左边
                img.draw_string(  0, 0, labels[i], color = 130, scale = 4)
            if (center_x >= 64):    # 中心坐标x >= 64,识别的图像在右边
                img.draw_string(100, 0, labels[i], color = 130, scale = 4)

            #img.draw_circle((center_x, center_y, 12), color=colors[i], thickness=2)

            lcd.display(img) # LCD显示识别到的数字，加上这句就会变卡
        #这里串口发送指令，
    print(clock.fps(), "fps", end="\n\n")
